"""
结果可视化模块：包含图像保存、探索性绘图、效应绘图与对比绘图等函数。
"""
from typing import List
import os
import yaml
import matplotlib.pyplot as plt
import numpy as np
import arviz as az
from PIL import Image


generated_images: List[str] = []


def save_fig(fig, filename, fig_save_path: str = None, plot_cfg: dict = None):
	"""Save a matplotlib figure using settings from plot_cfg (if provided).

	If fig_save_path is None, read `config.yaml` for `figures_dir` or `figure_path`.
	plot_cfg keys used: dpi, bbox_inches
	"""
	if plot_cfg is None:
		plot_cfg = {}
	dpi = plot_cfg.get("dpi", 300)
	bbox = plot_cfg.get("bbox_inches", "tight")

	# Determine figure save directory: explicit argument > config.yaml > default 'figures'
	if not fig_save_path:
		try:
			with open("config.yaml", "r", encoding="utf-8") as f:
				cfg = yaml.safe_load(f) or {}
		except Exception:
			cfg = {}
		fig_save_path = cfg.get("figures_dir") or cfg.get("figure_path") or "figures"

	os.makedirs(fig_save_path, exist_ok=True)
	path = os.path.join(fig_save_path, filename)
	fig.savefig(path, dpi=dpi, bbox_inches=bbox)
	generated_images.append(path)
	plt.close(fig)


def explore_and_save_figures(data, font, fig_save_path: str = None, plot_cfg: dict = None):
	"""Exploratory plots. plot_cfg keys used: figsize"""
	if plot_cfg is None:
		plot_cfg = {}
	figsize = tuple(plot_cfg.get('figsize', [15, 10]))
	plt.figure(figsize=figsize)
	plt.subplot(2, 2, 1)
	age_means = data.groupby('age')['income'].mean()
	plt.plot(age_means.index, age_means.values, marker='o')
	plt.title('年龄与收入关系', fontproperties=font)
	plt.xlabel('年龄', fontproperties=font)
	plt.ylabel('平均收入', fontproperties=font)

	plt.subplot(2, 2, 2)
	period_means = data.groupby('period')['income'].mean()
	plt.plot(period_means.index, period_means.values, marker='o')
	plt.title('时期与收入关系', fontproperties=font)
	plt.xlabel('时期', fontproperties=font)
	plt.ylabel('平均收入', fontproperties=font)

	plt.subplot(2, 2, 3)
	cohort_means = data.groupby('cohort')['income'].mean()
	plt.plot(cohort_means.index, cohort_means.values, marker='o')
	plt.title('队列与收入关系', fontproperties=font)
	plt.xlabel('队列 (出生年份)', fontproperties=font)
	plt.ylabel('平均收入', fontproperties=font)

	plt.subplot(2, 2, 4)
	plt.scatter(data['education'], data['income'], c=data['gender'], alpha=0.5)
	plt.title('教育与收入关系', fontproperties=font)
	plt.xlabel('教育年限', fontproperties=font)
	plt.ylabel('收入', fontproperties=font)

	plt.tight_layout()
	save_fig(plt.gcf(), "explore.png", fig_save_path, plot_cfg=plot_cfg)


def plot_and_save_trace(trace, title, filename, fig_save_path: str = None, plot_cfg: dict = None):
	"""Plot trace using arviz and save. plot_cfg keys may control dpi/bbox via save_fig."""
	if plot_cfg is None:
		plot_cfg = {}
	fig_array = az.plot_trace(trace)
	# az.plot_trace 返回一个 numpy 数组的 Figure/Axes 对象，取第一个 Figure
	try:
		fig = fig_array[0, 0].figure
	except Exception:
		# fallback: if it's a single figure
		fig = plt.gcf()
	plt.suptitle(title, fontsize=16)
	plt.tight_layout()
	save_fig(fig, filename, fig_save_path, plot_cfg=plot_cfg)


def plot_effect(trace, effect_name, n, categories, title, filename, font, fig_save_path: str = None, plot_cfg: dict = None):
	"""Plot effect estimates with HDI. plot_cfg passed to save_fig."""
	if plot_cfg is None:
		plot_cfg = {}
	# trace.posterior[effect_name].values shape: (chain, draw, n)
	samples3d = trace.posterior[effect_name].values
	mean = samples3d.mean(axis=(0, 1))
	# use arviz.hdi on array with shape (chain, draw, n) and specify dim for axis reduction
	# az.hdi will accept an array and compute hdi along the (chain, draw) axes per item
	hdi = az.hdi(samples3d, hdi_prob=0.95)
	# hdi shape: (n, 2) or xarray; normalize to ndarray with shape (n, 2)
	try:
		hdi_arr = hdi.values.reshape(n, 2)
	except Exception:
		hdi_arr = np.asarray(hdi)
	fig, ax = plt.subplots()
	ax.errorbar(
		x=range(n),
		y=mean,
		yerr=[mean - hdi_arr[:, 0], hdi_arr[:, 1] - mean],
		fmt='o', capsize=5
	)
	ax.axhline(y=0, color='r', linestyle='--')
	ax.set_title(title, fontproperties=font)
	ax.set_xlabel('组', fontproperties=font)
	ax.set_ylabel('效应大小', fontproperties=font)
	ax.set_xticks(range(n))
	ax.set_xticklabels([str(c) for c in categories], rotation=0, fontproperties=font)
	save_fig(fig, filename, fig_save_path, plot_cfg=plot_cfg)


def compare_effects(trace1, trace2, effect_name, n, categories, title, filename, font, fig_save_path: str = None, plot_cfg: dict = None, label1="APC模型", label2="完整模型"):
	"""Compare two effects. plot_cfg passed to save_fig."""
	if plot_cfg is None:
		plot_cfg = {}
	samples1_3d = trace1.posterior[effect_name].values
	mean1 = samples1_3d.mean(axis=(0, 1))
	hdi1 = az.hdi(samples1_3d, hdi_prob=0.95)
	try:
		hdi1_arr = hdi1.values.reshape(n, 2)
	except Exception:
		hdi1_arr = np.asarray(hdi1)

	samples2_3d = trace2.posterior[effect_name].values
	mean2 = samples2_3d.mean(axis=(0, 1))
	hdi2 = az.hdi(samples2_3d, hdi_prob=0.95)
	try:
		hdi2_arr = hdi2.values.reshape(n, 2)
	except Exception:
		hdi2_arr = np.asarray(hdi2)

	fig, ax = plt.subplots()
	ax.errorbar(range(n), mean1, yerr=[mean1 - hdi1_arr[:, 0], hdi1_arr[:, 1] - mean1], fmt='o-', capsize=4, label=label1)
	ax.errorbar(range(n), mean2, yerr=[mean2 - hdi2_arr[:, 0], hdi2_arr[:, 1] - mean2], fmt='s--', capsize=4, label=label2)
	ax.axhline(y=0, color='r', linestyle='--')
	ax.set_title(title, fontproperties=font)
	ax.set_xlabel('组', fontproperties=font)
	ax.set_ylabel('效应大小', fontproperties=font)
	ax.set_xticks(range(n))
	ax.set_xticklabels([str(c) for c in categories], rotation=0, fontproperties=font)
	ax.legend()
	save_fig(fig, filename, fig_save_path, plot_cfg=plot_cfg)


def show_all_images(generated_images: list):
	print("\n===== 统一弹窗显示所有图片 =====")
	for img_path in generated_images:
		try:
			img = Image.open(img_path)
			img.show()
		except Exception as e:
			print(f"图片显示失败: {img_path}，错误信息: {e}")
